from code.pretreat import get_small_samples
import os
import tensorflow as tf


# tfrecords文件的存储
def write_to_tfrecords(source_path,tfrecords_path):
    """
    将图片的特征值和目标值存进tfrecords
    :param source_path:源文件图片的路径
    :param tfrecords_path:生成tfrecords文件的路径
    :return: None
    """
    # 获取所有文件名
    file_names = os.listdir(source_path)
    file_names.remove('.gitkeep')
    for file_name in file_names:
        if file_name[-3:] == "xml":
            file_names.remove(file_name)

    # 记录当前是第几个大图
    step = 0
    # 截取路径的"test"或者"train"作为文件名
    test_or_train = tfrecords_path.split("\\")[-2]
    # 建立tfrecord存储器
    writer = tf.python_io.TFRecordWriter(tfrecords_path + "\\" + test_or_train+'0.tfrecords')
    for file_name in file_names:
        samples = get_small_samples(os.path.join(source_path,file_name))
        # 每8个大图存储到一个tfrecords文件中，文件序号0、1、2....n
        if step % 8 == 0:
            writer = tf.python_io.TFRecordWriter(tfrecords_path +test_or_train+'{}.tfrecords'.format(step // 8))
        step += 1
        # 循环将所有样本写入文件，每张图片样本都要构造example协议
        for sample in samples:
            # 构造一个样本的example
            example = tf.train.Example(features=tf.train.Features(feature={
                "image": tf.train.Feature(bytes_list=tf.train.BytesList(value=[sample.image.tostring()])),
                "flaw_id": tf.train.Feature(int64_list=tf.train.Int64List(value=[sample.flaw_id])),
            }))
            # 把单独的样本写入文件中
            writer.write(example.SerializeToString())
        print(file_name, " Done")
    # 关闭
    writer.close()


def read_from_tfrecords(tfrecords_path,batch_size):
    """
    tfrecords文件的读取
    :param tfrecords_path:tfrecords文件的路径
    :param batch_size:图片数量
    :return: images  flaw_ids
    """
    # 获取所有文件名
    file_names = os.listdir(tfrecords_path)
    file_paths = []
    # 获取文件的完整路径
    for file_name in file_names:
        file_paths.append(os.path.join(tfrecords_path,file_name))
    # 根据文件名生成一个文件队列
    filename_queue = tf.train.string_input_producer(file_paths)
    # 定义一个tfrecord的阅读器
    reader = tf.TFRecordReader()
    # 返回文件名和文件
    _, serialized_example = reader.read(filename_queue)
    # 取出包含image和flaw_id的feature对象
    features = tf.parse_single_example(serialized_example, features={
        'image': tf.FixedLenFeature([], tf.string),
        'flaw_id': tf.FixedLenFeature([], tf.int64),
    })
    # 将字符串解析成图片对应的像素组
    image = tf.decode_raw(features['image'], tf.uint8)
    # reshape为 224*224 的3通道图片，将图片还原成原来的维度
    image = tf.reshape(image, [224, 224, 3])
    # 在流中抛出flaw_id张量
    flaw_id = [tf.cast(features['flaw_id'], tf.int64)]

    image_batch, label_batch = tf.train.batch([image, flaw_id], batch_size=batch_size, num_threads=4, capacity=batch_size)
    return image_batch, label_batch


if __name__ == '__main__':
    source_train_path = os.path.abspath('../../data/train/')
    tfrecords_train_path = os.path.abspath('../../data/tfrecords/train/')
    # 建立train.tfrecords
    write_to_tfrecords(source_train_path, tfrecords_train_path)
    # 读取train.tfrecords
    train_images, train_flaw_ids = read_from_tfrecords(tfrecords_train_path, 32)
    # 开始一个会话
    with tf.Session() as sess:
        # 创建一个线程协调器，管理线程
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess, coord=coord)
        examples, ids = sess.run([train_images, train_flaw_ids])
        print(len(examples),len(ids),ids)
        coord.request_stop()
        # 等待线程完成
        coord.join(threads)
